{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from copy import deepcopy\n",
    "from time import time\n",
    "import os\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import tqdm\n",
    "from tqdm import tqdm\n",
    "import anndata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = \"./hest1k_datasets/PRAD/processed_data/\"  # processed data path\n",
    "ori_st_path = \"./hest1k_datasets/PRAD/st/\" # original ST data path\n",
    "\n",
    "slide_out = \"MEND145\"\n",
    "\n",
    "# load training slide IDs\n",
    "slidename_lst = list(np.genfromtxt(data_path + \"all_slide_lst.txt\", dtype=str))\n",
    "slidename_lst.remove(slide_out)\n",
    "print(slide_out, \" is held out for testing \", (slide_out not in slidename_lst))\n",
    "\n",
    "# load selected gene list\n",
    "selected_genes = list(np.genfromtxt(data_path + \"selected_gene_list.txt\", dtype=str))\n",
    "print(\"Selected gene list loaded. len of selected genes: \", len(selected_genes))\n",
    "\n",
    "# load test slide\n",
    "test_adata = anndata.read_h5ad(ori_st_path + slide_out + \".h5ad\")\n",
    "test_count_mtx_df = pd.DataFrame(test_adata.X.toarray(), columns=test_adata.var_names, index=test_adata.obs_names).loc[:, selected_genes]\n",
    "# transform count\n",
    "test_count_mtx_selected_genes = np.log2(test_count_mtx_df + 1).copy()\n",
    "print(\"Test count mtx shape: \", test_count_mtx_selected_genes.shape)\n",
    "\n",
    "# load generated samples for test slide\n",
    "sample_path = \"./PRAD_results/runs/000/samples/generated_samples_0300000_20sample.pt\" # (example)\n",
    "pred = torch.load(sample_path, )\n",
    "pred = pred.squeeze(1)\n",
    "print(\"Generated samples shape: \", pred.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_rep = 20       # number of samples generated for one image patch\n",
    "num_selected = 20  # take average of (a subset / all) samples and use it as prediction\n",
    "random_selected_index = np.random.choice(np.arange(num_rep), num_selected, replace=False)\n",
    "\n",
    "pred_avg = torch.zeros(size=(test_count_mtx_selected_genes.shape[0], \n",
    "                             test_count_mtx_selected_genes.shape[1]))\n",
    "for i in range(test_count_mtx_selected_genes.shape[0]):\n",
    "    pred_avg[i] = torch.mean(pred[i*num_rep + random_selected_index, :], dim=0)\n",
    "\n",
    "pred_avg = pred_avg.cpu().detach().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate correlation\n",
    "\n",
    "all_corr = []\n",
    "for i in range(test_count_mtx_selected_genes.shape[1]):\n",
    "    x = test_count_mtx_selected_genes.iloc[:, i].values\n",
    "    y = pred_avg[:, i]\n",
    "    cor = np.corrcoef(x, y)[0][1]\n",
    "    all_corr.append(cor)\n",
    "plt.hist(all_corr, bins=50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluation metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# evaluation metrics\n",
    "\n",
    "# PCC\n",
    "print(\"PCC-10: \", np.mean(sorted(all_corr)[::-1][:10]))\n",
    "print(\"PCC-50: \", np.mean(sorted(all_corr)[::-1][:50]))\n",
    "print(\"PCC-200: \", np.mean(sorted(all_corr)[::-1][:200]))\n",
    "# MSE, MAE\n",
    "print(\"MSE: \", np.mean((test_count_mtx_selected_genes.values - pred_avg)**2))\n",
    "print(\"MAE: \", np.mean(np.abs(test_count_mtx_selected_genes.values - pred_avg)))\n",
    "# RVD\n",
    "pred_var = np.var(pred_avg, axis=0)\n",
    "gt_var = np.var(test_count_mtx_selected_genes.values, axis=0)\n",
    "print(\"RVD: \", np.mean((pred_var - gt_var)**2 / gt_var**2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# gene variation curve\n",
    "fig, axs = plt.subplots(2, 2, figsize=(8, 8))\n",
    "\n",
    "pred_mean = np.mean(pred_avg, axis=0)\n",
    "pred_mean = pred_mean / np.sum(pred_mean)\n",
    "gt_mean = np.mean(test_count_mtx_selected_genes, axis=0)\n",
    "gt_mean = gt_mean / np.sum(gt_mean)\n",
    "gt_mean_sorted = np.sort(gt_mean)\n",
    "pred_mean_sorted = pred_mean[np.argsort(gt_mean)]\n",
    "axs[0, 0].plot(np.arange(len(gt_mean_sorted)), gt_mean_sorted, label=\"Ground Truth\", c=\"b\")\n",
    "axs[0, 0].scatter(np.arange(len(pred_mean_sorted)), pred_mean_sorted, s=5, label=\"Predicted\", c=\"orange\")\n",
    "axs[0, 0].set_title(\"Normalized Mean\")\n",
    "axs[0, 0].set_xlabel(\"gene index ordered by mean\")\n",
    "axs[0, 0].set_ylabel(\"normalized mean\")\n",
    "# axs[0, 0].set_ylim()\n",
    "\n",
    "pred_mean = np.mean(pred_avg, axis=0)\n",
    "gt_mean = np.mean(test_count_mtx_selected_genes, axis=0)\n",
    "gt_mean_sorted = np.sort(gt_mean)\n",
    "pred_mean_sorted = pred_mean[np.argsort(gt_mean)]\n",
    "axs[1, 0].plot(np.arange(len(gt_mean_sorted)), gt_mean_sorted, label=\"Ground Truth\", c=\"b\")\n",
    "axs[1, 0].scatter(np.arange(len(pred_mean_sorted)), pred_mean_sorted, s=5, label=\"Predicted\", c=\"orange\")\n",
    "axs[1, 0].set_title(\"Absolute Mean\")\n",
    "axs[1, 0].set_xlabel(\"gene index ordered by mean\")\n",
    "axs[1, 0].set_ylabel(\"absolute mean\")\n",
    "# axs[1, 0].set_ylim()\n",
    "\n",
    "\n",
    "pred_var = np.var(pred_avg, axis=0)\n",
    "pred_var = pred_var / np.sum(pred_var)\n",
    "gt_var = np.var(test_count_mtx_selected_genes, axis=0)\n",
    "gt_var = gt_var / np.sum(gt_var)\n",
    "gt_var_sorted = np.sort(gt_var)\n",
    "pred_var_sorted = pred_var[np.argsort(gt_var)]\n",
    "axs[0, 1].plot(np.arange(len(gt_var_sorted)), gt_var_sorted, label=\"Ground Truth\", c=\"b\")\n",
    "axs[0, 1].scatter(np.arange(len(pred_var_sorted)), pred_var_sorted, s=5, label=\"Predicted\", c=\"orange\")\n",
    "axs[0, 1].set_title(\"Normalized Variance\")\n",
    "axs[0, 1].set_xlabel(\"gene index ordered by var\")\n",
    "axs[0, 1].set_ylabel(\"normalized variance\")\n",
    "# axs[0, 1].set_ylim()\n",
    "\n",
    "\n",
    "\n",
    "pred_var = np.var(pred_avg, axis=0)\n",
    "gt_var = np.var(test_count_mtx_selected_genes, axis=0)\n",
    "gt_var_sorted = np.sort(gt_var)\n",
    "pred_var_sorted = pred_var[np.argsort(gt_var)]\n",
    "axs[1, 1].plot(np.arange(len(gt_var_sorted)), gt_var_sorted, label=\"Ground Truth\", c=\"b\")\n",
    "axs[1, 1].scatter(np.arange(len(pred_var_sorted)), pred_var_sorted, s=5, label=\"Predicted\", c=\"orange\")\n",
    "axs[1, 1].set_title(\"Absolute Variance\")\n",
    "axs[1, 1].set_xlabel(\"gene index ordered by var\")\n",
    "axs[1, 1].set_ylabel(\"absolute variance\")\n",
    "# axs[1, 1].set_ylim()\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot marker genes\n",
    "\n",
    "Image.MAX_IMAGE_PIXELS = None\n",
    "\n",
    "img_path = \"./hest1k_datasets/PRAD/wsis/\"\n",
    "img_raw = Image.open(img_path + slide_out + \".tif\") # change suffix according to image type\n",
    "\n",
    "x = test_adata.obsm[\"spatial\"][:, 0]\n",
    "y = test_adata.obsm[\"spatial\"][:, 1]\n",
    "for gene in [\"\"]: # fill in gene names\n",
    "    fig, axs = plt.subplots(1, 2, figsize=(8, 4))\n",
    "    fig.suptitle(gene)\n",
    "    gene_idx = np.where(np.array(selected_genes) == gene)[0][0]\n",
    "\n",
    "    axs[0].imshow(img_raw)\n",
    "    color_gt = test_adata[:, test_adata.var_names == gene].X.toarray().flatten()\n",
    "    im0 = axs[0].scatter(x, y, c=color_gt, s=15, alpha=0.9)\n",
    "    axs[0].set_title(\"ground truth\")\n",
    "    fig.colorbar(im0)\n",
    "    \n",
    "    axs[1].imshow(img_raw)\n",
    "    color_pred = pred_avg[:, gene_idx]\n",
    "    im1 = axs[1].scatter(x, y, c=color_pred, s=15, alpha=0.9) \n",
    "    axs[1].set_title(\"pred\")\n",
    "    fig.colorbar(im1)\n",
    "    \n",
    "    plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "DiT",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
